import csv

import scipy.io as scio
import numpy as np
import os
import torch
from CR_DSPPytorch import FIRLayer, PhaseRecLayer, EDCLayer
import util

if __name__ == '__main__':
    BASE_DIR = os.path.dirname(__file__)
    trSetDir = os.path.join(BASE_DIR, 'data/dataset')
    lpPick = [2]
    trSetDirSNRVaryDir = os.path.join(trSetDir, 'SNRVary')

    if not os.path.exists(trSetDirSNRVaryDir):
        os.makedirs(trSetDirSNRVaryDir)

    for lpIndx, lp in enumerate(lpPick):
        data = scio.loadmat(os.path.join(trSetDir, f'trSet_dBm_{lp}_nn_sf_1.mat'))
        trsetMat = data['dataset']
        trsetMat = np.transpose(trsetMat, (1, 0))
        prbs = np.stack([data['prbsx'].squeeze(), data['prbsy'].squeeze()], axis=0)
        label = np.stack([data['labelx'].squeeze(), data['labely'].squeeze()], axis=0)

        awgn = np.random.randn(*trsetMat.shape)
        SNR_range = [5, 7, 9, 11, 13, 15, 17, 19]

        berBaseline = util.pr_ber(trsetMat, prbs, util.CONST_16QAM)[1]
        QBaseline = util.ber2q(np.mean(berBaseline))
        print(f"Linear Equalization Result, Q: {QBaseline}")

        Q_corresp = []

        for SNR in SNR_range:
            dict_corresponding = data.copy()
            noise_std = 1 / np.sqrt(2 * 10.0 ** (SNR / 10.0))
            trsetMatAWGN = trsetMat + awgn * noise_std
            ber = util.pr_ber(trsetMatAWGN, prbs)[1]
            print(f"The input Q factor is : {util.ber2q(np.mean(ber))}")
            Q_corresp.append(util.ber2q(np.mean(ber)))
            dict_corresponding['dataset'] = trsetMatAWGN.T
            trSetDirSNRVaryMatPath = os.path.join(trSetDirSNRVaryDir, f'trset_add_AWGN_by_SNR_{SNR}.mat')
            scio.savemat(trSetDirSNRVaryMatPath, dict_corresponding)
            del dict_corresponding

        with open(os.path.join(trSetDirSNRVaryDir, f'input_Q_result_corresponding_SNR_at_{lp}_dBm'), 'w') as f:
            csv_writer = csv.writer(f)
            csv_writer.writerow(['SNR', 'Input Q factor'])
            csv_writer.writerow(['LE result', QBaseline])
            for indx, q in enumerate(Q_corresp):
                csv_writer.writerow([SNR_range[indx], Q_corresp[indx]])

